import random
import math
import numpy as np
from PIL import Image


class MaskingGenerator:
    def __init__(
            self, input_size, num_masking_patches, min_num_patches=4, max_num_patches=None,
            min_aspect=0.3, max_aspect=None):
        if not isinstance(input_size, tuple):
            input_size = (input_size, ) * 2
        self.height, self.width = input_size

        self.num_patches = self.height * self.width
        self.num_masking_patches = num_masking_patches

        self.min_num_patches = min_num_patches
        self.max_num_patches = num_masking_patches if max_num_patches is None else max_num_patches

        max_aspect = max_aspect or 1 / min_aspect
        self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect))

    def __repr__(self):
        repr_str = "Generator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % (
            self.height, self.width, self.min_num_patches, self.max_num_patches,
            self.num_masking_patches, self.log_aspect_ratio[0], self.log_aspect_ratio[1])
        return repr_str

    def get_shape(self):
        return self.height, self.width

    def _mask(self, mask, max_mask_patches):
        delta = 0
        for attempt in range(10):
            target_area = random.uniform(self.min_num_patches, max_mask_patches)
            aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
            h = int(round(math.sqrt(target_area * aspect_ratio)))
            w = int(round(math.sqrt(target_area / aspect_ratio)))
            if w < self.width and h < self.height:
                top = random.randint(0, self.height - h)
                left = random.randint(0, self.width - w)

                num_masked = mask[top: top + h, left: left + w].sum()
                # Overlap
                if 0 < h * w - num_masked <= max_mask_patches:
                    for i in range(top, top + h):
                        for j in range(left, left + w):
                            if mask[i, j] == 0:
                                mask[i, j] = 1
                                delta += 1

                if delta > 0:
                    break
        return delta

    def __call__(self):
        mask = np.zeros(shape=self.get_shape(), dtype=int)
        mask_count = 0
        while mask_count != self.num_masking_patches:
            max_mask_patches = self.num_masking_patches - mask_count
            max_mask_patches = min(max_mask_patches, self.max_num_patches)

            delta = self._mask(mask, max_mask_patches)
            mask_count += delta
        
        return mask


class RandomMaskingGenerator:
    def __init__(self, input_size, mask_ratio):
        if not isinstance(input_size, tuple):
            input_size = (input_size,) * 2

        self.height, self.width = input_size

        self.num_patches = self.height * self.width
        self.num_mask = int(mask_ratio * self.num_patches)

    def __repr__(self):
        repr_str = "Maks: total patches {}, mask patches {}".format(
            self.num_patches, self.num_mask
        )
        return repr_str

    def __call__(self):
        mask = np.hstack([
            np.zeros(self.num_patches - self.num_mask),
            np.ones(self.num_mask),
        ])
        np.random.shuffle(mask)
        return mask # [196]


class MaskGenerator(object):
    def __init__(self, image_size=(224, 224), mask_ratio: float = 0.5, window_size: int = 16):
        super(MaskGenerator, self).__init__()
        if not isinstance(image_size, tuple):
            image_size = (image_size, image_size)
        self.h, self.w = image_size
        assert image_size[0] % window_size == 0 and image_size[1] % window_size == 0, "Not divisible"

        self.window_size=window_size
        self.h_patch = self.h // window_size
        self.w_patch = self.w // window_size
        self.mask_ratio = mask_ratio
        self.num_patches = self.h_patch * self.w_patch
        self.mask_patchs = int(self.mask_ratio * self.num_patches)
        self.patch_w = window_size
        self.patch_h = window_size
        self.idx = np.hstack([
            np.zeros(self.mask_patchs),
            np.ones(self.num_patches - self.mask_patchs),
        ])
   
        self.mask_img_hw=np.array(self.crack(self.mask_patchs))*window_size
        self.visible_img_hw=np.array(self.crack(self.num_patches-self.mask_patchs))*window_size

    def crack(self,integer):
        start = int(np.sqrt(integer))
        factor = integer / start
        while not self.is_integer(factor):
            start += 1
            factor = integer / start
        return int(factor), start


    def is_integer(self,number):
        if int(number) == number:
            return True
        else:
            return False

    def __call__(self, img):
        idx = np.random.permutation(self.idx)
        mask = np.zeros((self.h, self.w, 1))
        visible_list=[]
        mask_list=[]
        channel=img.shape[2]
        cur = 0
        for i in range(0, self.h, self.window_size):
            for j in range(0, self.w, self.window_size):
                mask[i:i+self.window_size, j:j+self.window_size, :] = idx[cur]
                if idx[cur]==0:
                    mask_list.append(img[i:i+self.window_size, j:j+self.window_size, :])
                else:
                    visible_list.append(img[i:i+self.window_size, j:j+self.window_size, :])
                cur += 1
        
        visible_img=np.zeros((self.visible_img_hw[0],self.visible_img_hw[1],channel))
        
        cur=0
        for i in range(0,self.visible_img_hw[0],self.window_size):
            for j in range(0,self.visible_img_hw[1],self.window_size):
                visible_img[i:i+self.window_size,j:j+self.window_size,:]=visible_list[cur]
                cur+=1

        mask_img=np.zeros((self.mask_img_hw[0],self.mask_img_hw[1],channel))
        cur=0
        for i in range(0,self.mask_img_hw[0],self.window_size):
            for j in range(0,self.mask_img_hw[1],self.window_size):
                mask_img[i:i+self.window_size,j:j+self.window_size,:]=mask_list[cur]
                cur+=1
            
        return img * mask,visible_img,mask_img


if __name__ == "__main__":
    
    mask_generator = MaskGenerator(image_size= 224, mask_ratio=0.75, window_size=4)
    img = np.array(Image.open(
        r'/home/ubuntu/code/pytorch_code/data/ADE_train_00008223.jpg').resize(size=(224, 224)))
    print(img.shape)
    mask,visible_img,mask_img = mask_generator(img)
    print("mask",mask.shape)
    print("visible_img",visible_img.shape)
    print("mask_img",mask_img.shape)
    mask[mask==0]=128
    img = Image.fromarray(np.uint8(mask))
    img.show()
    img.save(r"/home/ubuntu/code/pytorch_code/test/1.png")

    visible_img = Image.fromarray(np.uint8(visible_img))
    visible_img.show()
    visible_img.save(r"/home/ubuntu/code/pytorch_code/test/visible_img.png")

    mask_img = Image.fromarray(np.uint8(mask_img))
    mask_img.show()
    mask_img.save(r"/home/ubuntu/code/pytorch_code/test/mask_img.png")






